package com.shujia.mllib

import org.apache.spark.sql.SparkSession
import org.apache.spark.ml.regression.{LinearRegression, LinearRegressionModel}
import org.apache.spark.ml.feature.{OneHotEncoder, StandardScaler,
  StringIndexer, VectorAssembler}
import org.apache.spark.ml.{Pipeline, PipelineModel}
import org.apache.spark.sql.types._
import org.apache.spark.sql.functions._
import org.apache.spark.ml.evaluation.RegressionEvaluator
object Test4 {
  def main(args: Array[String]): Unit = {
    // ===================================================================
    // 1. Spark环境初始化
    // ===================================================================
    val spark = SparkSession.builder()
      .appName("HousePricePrediction") // 定义应用名称（显示在Spark UI）
      .master("local[*]") // 本地模式运行，使用所有CPU核心
      .getOrCreate() // 获取或创建会话（防止重复创建）
    import spark.implicits._ // 启用DataFrame隐式转换（如RDD转DF）
    // ===================================================================
    // 2. 数据加载与验证
    // ===================================================================
    // 定义数据模式（明确字段类型，避免自动推断错误）
    val schema = StructType(Array(
      StructField("area", IntegerType), // 数值特征：面积（单位平方米）
      StructField("bedrooms", IntegerType), // 数值特征：卧室数量
      StructField("bathrooms", IntegerType), // 数值特征：卫生间数量
      StructField("district", StringType), // 分类特征：区域（需编码处理）
      StructField("age", IntegerType), // 数值特征：房龄（单位年）
      StructField("floor", IntegerType), // 数值特征：楼层
      StructField("price", DoubleType) // 目标变量：房价（单位万元）
    ))
    // 读取CSV文件（指定模式避免类型错误）
    val rawData = spark.read
      .option("header", "true") // 第一行为列名（自动对齐schema）
      .schema(schema) // 应用预定义模式（确保数据一致性）
      .csv("spark/data/house_prices.csv") // 数据存储路径
    // ===================================================================
    // 3. 特征工程管道构建
    // ===================================================================
    // 3.1 区域特征编码（处理分类变量）
    val districtIndexer = new StringIndexer()
      .setInputCol("district") // 原始字符串类型列（A/B/C/D）
      .setOutputCol("district_index") // 转换为数值索引（如A→0.0, B→1.0）
    val districtEncoder = new OneHotEncoder()
      .setInputCol("district_index") // 输入数值索引列
      .setOutputCol("district_vec") // 输出独热编码向量（稀疏向量表示）
    // 3.2 特征向量组装（合并所有特征）
    val assembler = new VectorAssembler()
      .setInputCols(Array( // 选择需要合并的特征列
        "area", "bedrooms", "bathrooms",
        "district_vec", "age", "floor"
      ))
      .setOutputCol("raw_features") // 输出未标准化的密集向量
    // 3.3 特征标准化（消除量纲影响）
    val scaler = new StandardScaler()
      .setInputCol("raw_features") // 输入组装后的特征向量
      .setOutputCol("features") // 输出标准化后的特征
      .setWithStd(true) // 启用标准差缩放（使特征方差为1）
      .setWithMean(true) // 启用均值居中（使特征均值为0）
    // ===================================================================
    // 4. 机器学习模型构建
    // ===================================================================
    val lr = new LinearRegression()
      .setLabelCol("price") // 设置目标变量（要预测的值）
      .setFeaturesCol("features") // 设置特征向量列
      .setMaxIter(100) // 设置最大迭代次数（确保收敛）
      .setRegParam(0.3) // 设置L2正则化参数（防止过拟合）
    // 构建处理管道（按顺序执行数据处理步骤）
    val pipeline = new Pipeline()
      .setStages(Array( // 定义管道阶段执行顺序：
        districtIndexer, districtEncoder, // 先处理分类特征
        assembler, scaler, // 再合并和标准化特征
        lr // 最后训练模型
      ))
    // ===================================================================
    // 5. 模型训练与验证
    // ===================================================================
    // 拆分数据集（保持数据分布一致性）
    val Array(trainData, testData) = rawData.randomSplit(Array(0.8, 0.2), seed = 42)
    // 90%训练集用于模型训练，10%测试集用于验证，seed固定保证可复现性
    // 训练模型（管道会自动执行所有预处理步骤）
    val model = pipeline.fit(trainData) // 输出包含预处理步骤和训练好的模型
    // 进行预测（自动应用相同的预处理流程）
    val predictions = model.transform(testData)
      .select("price", "prediction") // 选择真实值和预测值用于评估
    // ===================================================================
    // 6. 模型评估与解释
    // ===================================================================
    // 初始化评估器（统一评估框架）
    val evaluator = new RegressionEvaluator()
      .setLabelCol("price") // 设置真实值列
      .setPredictionCol("prediction") // 设置预测值列
    // 计算评估指标（多维度评估模型性能）
    val rmse = evaluator.setMetricName("rmse").evaluate(predictions) // 均方根误差（对异常值敏感）
    val mae = evaluator.setMetricName("mae").evaluate(predictions) // 平均绝对误差（更易解释）
    val r2 = evaluator.setMetricName("r2").evaluate(predictions) // R平方值（模型解释能力）
    // 获取模型细节（理解特征影响）
    val lrModel = model.stages.last.asInstanceOf[LinearRegressionModel]
    println(
      s"""模型评估报告：
         |================================
         |RMSE(均方根误差): $rmse → 预测值与真实值的平均偏离程度，单位与房价相同
         |MAE(平均绝对误差): $mae → 平均每个预测误差的绝对值
         |R²(决定系数): $r2 → 模型解释了${"%.1f".format(r2*100)}%%的房价方差
         |
         |模型参数解读：
         |截距项: ${lrModel.intercept} → 所有特征为0时的基准房价（需结合标准化理解）
         |特征权重:
         | - area: ${lrModel.coefficients(0)} → 面积每增加1平米，房价变化量
         | - bedrooms: ${lrModel.coefficients(1)} → 卧室数量对房价的影响
         | - bathrooms: ${lrModel.coefficients(2)} → 卫生间数量的边际效应
         | - district_vec: ${lrModel.coefficients(3)} → 区域相对基准的溢价系数
         | - age: ${lrModel.coefficients(4)} → 房龄每增加1年的影响
         | - floor: ${lrModel.coefficients(5)} → 楼层高度的价值系数
         |================================
         |""".stripMargin)
    // ===================================================================
    // 7. 交互式预测（业务应用）
    // ===================================================================
    println("\n=== 房价预测系统（输入exit退出） ===")
    while (true) {
      print("输入格式：面积,卧室数,卫生间数,区域(A/B/C/D),房龄,楼层\n> ")
      val input = scala.io.StdIn.readLine().trim
      if (input.equalsIgnoreCase("exit")) System.exit(0)
      try {
        // 解析用户输入（格式校验）
        val Array(area, bedrooms, bathrooms, district, age, floor) =
          input.split(",")
        // 创建临时DataFrame（必须包含所有管道需要的字段）
        val df = Seq((area.toInt, bedrooms.toInt, bathrooms.toInt,
          district, age.toInt, floor.toInt, 0.0)) // price设为0（模型不需要真实值）
          .toDF("area","bedrooms","bathrooms","district","age","floor", "price")
        // 生成预测（自动应用管道预处理）
        val prediction = model.transform(df)
          .select("prediction") // 提取预测结果
          .first() // 获取第一条（唯一）记录
          .getDouble(0) // 转换为Double类型
        // 格式化输出（保留1位小数）
        println(f"预测房价: ￥${prediction}%.1f万元")
        println("------------------------------")
      } catch {
        case e: Exception =>
          println("错误！请按格式输入，示例：120,3,2,A,5,10")
      }
    }
  }
}
